/**
 * 
 */
package timingAttack;

import java.io.Console;
import java.io.FileInputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStream;
import java.io.PrintWriter;
import java.math.BigInteger;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.RSAPrivateKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.util.Arrays;
import java.util.Random;

/**
 * @author Vincent
 *
 */
public class SimpleDemoMean
{
	// For debugging
	static final boolean verbose = true;

	// Constants for the parameters of the demo
	static final int sampleSize = 10000;	// number of messages to try
	static final int modulusSize = 4096;	// size of modulus in bits
	static final int messageSize = 3095;	// size of messages in bits

	static BigInteger modulus;
	static BigInteger exponent; 
	static BigInteger guessForExponent;
	
	// These are the messages for which we will time modexp operations
	static BigInteger [] eMessage = new BigInteger[sampleSize];
	static BigInteger [] fMessage = new BigInteger[sampleSize];

	// These are the total modular exponentiation times
	static long [] eModExpTime; 
	static long [] fModExpTime; 

	// These are the total number of messages of the two types
	static int eMessageCount;
	static int fMessageCount;
	
	// Flags for file output
	static final boolean modExpFileOutput = false;
	static final boolean bitFileOutput = false;
	static final boolean modulusFileOutput = true;

	// Filenames for file output
	static String modExpTimeFilename = modulusSize + "." + messageSize + ".modExpTime";
	static String unsetBitFilename = modulusSize + "." + messageSize + ".unsetBit";
	static String setBitFilename = modulusSize + "." + messageSize + ".setBit";
	static String modulusFilename = modulusSize + ".modulus";

	
	/**
	 * This is the timing method to use
	 * Separating it like this makes it easier to
	 * switch between using System.nanoTime() and
	 * System.currentTimeMillis()
	 * 
	 * @return timing value
	 */
	static long time()
	{
		// return System.currentTimeMillis();
		return System.nanoTime();
	}
	
	/**
	 * @param args
	 */
	public static void main(String[] args) 
	{
		
		// SecureRandom r = new SecureRandom();
		// exponent = new BigInteger(256,r);
		
		// Console console = System.console();
		// new BigInteger(console.readLine("Enter the exponent to guess: "));
		
		exponent = new BigInteger("101"); 
		
		System.out.println("Exponent decimal representation: " + exponent.toString(10));
		System.out.println("Exponent binary representation: " + exponent.toString(2));
		System.out.println("Number of bits to guess: " + exponent.bitLength());

		// Obtain modulus
		System.out.print("Generating " + modulusSize + "-bit modulus...");
		try
		{
			modulus = generateModulus(modulusSize);
		}
		catch (Exception e)
		{
			e.printStackTrace();
			return;
		}
		
		// Save modulus to file
		if( modulusFileOutput )
		{
			PrintWriter p;
			try 
			{
				p = new PrintWriter(new FileWriter(modulusFilename));
			} 
			catch (IOException e) 
			{
				e.printStackTrace();
				return;
			}
			p.print(modulus);
			p.close();
		}
		
		// Generate timing data for messages
		System.out.print("\nGenerating timing data for " + sampleSize + " " + messageSize + "-bit messages...");
		timeModExp(sampleSize,messageSize,exponent,modulus);

		// Guess the bits
		System.out.println("\n\nGuessing bits..");
		guessForExponent = guessBits(sampleSize);

		// Print out answer
		System.out.println("\nExponent binary representation: " + exponent.toString(2));
		System.out.println("Guess binary representation: " + guessForExponent.toString(2));
	}

	/**
	 * Print all the elements of an array to a file
	 * 
	 * @param data
	 * @param p
	 */
	static void writeArrayToFile( long [] data, PrintWriter p )
	{
		for( int i = 0; i < data.length; i++ )
		{
			p.println(data[i]);
		}
	}
	
	/***
	 * Generate a modulus of a specified size
	 * @throws NoSuchAlgorithmException 
	 * @throws InvalidKeySpecException 
	 * 
	 */
	public static BigInteger generateModulus(int size) throws NoSuchAlgorithmException, InvalidKeySpecException
	{
		// Save some time and try to read modulus from file
		try
		{
			InputStream in = new FileInputStream(modulusFilename);
			byte [] modulusBytes = new byte[in.available()];
			in.read( modulusBytes );
			in.close();
			// System.out.println("Length = " + modulusBytes.length);
			// System.out.println("Number =" + new String(modulusBytes,"UTF-8").substring(modulusBytes.length-10));
			BigInteger b = new BigInteger( new String(modulusBytes,"UTF-8") );
			System.out.print(b.bitLength() + "-bit modulus read from " + modulusFilename);
			return b;
		}
		catch(Exception e)
		{
		}
		
		// Cannot obtain modulus from file--create a new one from an RSA public key
		// http://www.javamex.com/tutorials/cryptography/rsa_encryption.shtml
		KeyPairGenerator kpg = KeyPairGenerator.getInstance("RSA");
		kpg.initialize(size);
		KeyPair kp = kpg.genKeyPair();
		Key publicKey = kp.getPublic();
		Key privateKey = kp.getPrivate();
		KeyFactory fact = KeyFactory.getInstance("RSA");
		RSAPublicKeySpec pub = fact.getKeySpec(kp.getPublic(), RSAPublicKeySpec.class);
		RSAPrivateKeySpec priv = fact.getKeySpec(kp.getPrivate(), RSAPrivateKeySpec.class);

		return pub.getModulus();
	}
	
	/***
	 * Generate timing data for the modular exponentiation of a
	 * specified number of messages of a given size
	 * 
	 * @param messageCount number of messages for timing
	 * @param size message size
	 * @param messageArray array to store messages in
	 * @param e exponent for modular exponentiation
	 * @param m modulus for modular exponentiation
	 * @return
	 */
	public static void timeModExp(int messageCount, int size, BigInteger e, BigInteger m)
	{
		long [] time = new long[messageCount];
		Random r = new SecureRandom();
		long start = 0, end = 0;
		eModExpTime = new long[messageCount];
		fModExpTime = new long[messageCount];
		
		for( int i = 0; i < messageCount/2; i++ )
		{
			eMessage[i] = getE(modulusSize/3,r);
			fMessage[i] = getF(modulusSize/2,r);
			
			// Time modular exponentiation of message
            start = time();
        	ModExp.modExp( eMessage[i], e, m );
            end = time();

			eModExpTime[i] = (end - start);
			
			// Time modular exponentiation of message
            start = time();
        	ModExp.modExp( fMessage[i], e, m );
            end = time();

			fModExpTime[i] = (end - start);

			if ((i % (messageCount/10)) == 0)
				System.out.print(i*100/messageCount + "%...");
		}
		eMessageCount = messageCount/2;
		fMessageCount = messageCount/2;
		
		System.out.println();
	}

	static BigInteger getE(int size, Random r)
	{
		boolean foundOne = false;
		BigInteger candidate = null;
		while( !foundOne )
		{
			candidate = new BigInteger(size,r);
			foundOne = (candidate.pow(3).compareTo(modulus) < 0 );
			// System.out.println(foundOne + "E");
		}
		
		return candidate;
	}
	
	static BigInteger getF(int size, Random r)
	{
		boolean foundOne = false;
		BigInteger candidate = null;
		while( !foundOne )
		{
			candidate = new BigInteger(size,r);
			foundOne = (candidate.pow(3).compareTo(modulus) > 0 ) && (candidate.pow(2).compareTo(modulus) < 0);
			// System.out.println(foundOne + "F");
		}
		
		return candidate;
	}
	
	static boolean isEValue(BigInteger b)
	{
		return (b.pow(3).compareTo(modulus) < 0 );
	}
	
	static boolean isFValue(BigInteger b)
	{
		return (b.pow(3).compareTo(modulus) > 0 ) && (b.pow(2).compareTo(modulus) < 0);
	}
	
	static boolean fLarger( long [] eTimes, int eCount, long [] fTimes, int fCount )
	{
		long [] e = Arrays.copyOf(eTimes, eCount);
		long [] f = Arrays.copyOf(fTimes, fCount);

		System.out.println("f-bar = " + Statistics.mean(f) + " e-bar = " + Statistics.mean(e));
		return Statistics.mean(f) > Statistics.mean(e);
	}
	
	
	
	/**
	 * Guess the bits of a modular exponentiation exponent using timing data
	 * 
	 * @param messageCount number timed modular exponentation operations
	 * @return
	 */
	static BigInteger guessBits(int messageCount)
	{
		double eMean = Statistics.mean(eModExpTime);
		double fMean = Statistics.mean(fModExpTime);

		System.out.println(eMean + " " + fMean);

		BigInteger [] eIntermediateValues = new BigInteger [messageCount/2];
		BigInteger [] fIntermediateValues = new BigInteger [messageCount/2];

		long [] newETimes = new long[messageCount];
		long [] newFTimes = new long[messageCount];
		int eCount = 0;
		int fCount = 0;		
		
		for( int i = 0; i < messageCount/2; i++ )
		{
			eIntermediateValues[i] = eMessage[i].pow(3).mod(modulus);
			fIntermediateValues[i] = fMessage[i].pow(3).mod(modulus);

			if( isEValue(eIntermediateValues[i]) )
				newETimes[eCount++] = eModExpTime[i];
			else if( isFValue(eIntermediateValues[i]) )
				newFTimes[fCount++] = eModExpTime[i];
				
			if( isEValue(fIntermediateValues[i]) )
				newETimes[eCount++] = fModExpTime[i];
			else if( isFValue(fIntermediateValues[i]) )
				newFTimes[fCount++] = fModExpTime[i];		
		}
		
		System.out.println("eCount = " + eCount + " fCount = " + fCount);
		// fLarger(newETimes,eCount,newFTimes,fCount);
		
//		for( int i = 0; i < sampleSize/2; i++ )
//		{
//			System.out.println("e = " + eModExpTime[i] + " f = " + fModExpTime[i]);
//		}
		return null;
		
		
		
//		BigInteger guess = new BigInteger("0");
//		BigInteger [] runningTotal = new BigInteger[messageCount];	// of modular exp operation
//		BigInteger [][] intermediateValue = 							// candidate values
//			{ new BigInteger [messageCount], new BigInteger [messageCount] }; 
//
//		long [][] remainingTime = { new long[messageCount], new long[messageCount] };
//		long start = 0, end = 0, squareTime = 0, multTime = 0;
//		
//		// Initialize running totals, iteration times
//		for( int i = 0; i < messageCount; i++ )
//		{
//			remainingTime[0][i] = remainingTime[1][i] = modExpTime[i];
//			runningTotal[i] = BigInteger.ONE;
//		}
//
//		// Make a guess for each bit of the exponent
//		for( int i = 0; i < exponent.bitLength(); i++ )
//		{
//			// Open files to record times for current bit
//			PrintWriter p0, p1;
//			if( bitFileOutput )
//			{
//				try 
//				{
//					p0 = new PrintWriter(new FileWriter(unsetBitFilename+Integer.toString(i)));
//					p1 = new PrintWriter(new FileWriter(setBitFilename+Integer.toString(i)));
//				} 
//				catch (IOException e) 
//				{
//					e.printStackTrace();
//					return null;
//				}
//			}
//			
//			// Calculate times for two possibilities for each message
//			for( int j = 0; j < messageCount; j++ )
//			{
//				// time for bit not set
//                start = time();
//				intermediateValue[0][j] = runningTotal[j].multiply(runningTotal[j]);
//				if(intermediateValue[0][j].compareTo(modulus) > 0)
//					intermediateValue[0][j] = intermediateValue[0][j].mod(modulus);
//				// options[0][j] = runningTotal[j].multiply(runningTotal[j]).mod(modulus);
//                end = time();
//
//				squareTime = (end - start);
//				remainingTime[0][j] -= squareTime;
//				if( bitFileOutput ) p0.println(squareTime);
//				
//				// time for bit set
//                start = time();
//				intermediateValue[1][j] = intermediateValue[0][j].multiply(message[j]);
//				if(intermediateValue[1][j].compareTo(modulus) > 0)
//					intermediateValue[1][j] = intermediateValue[1][j].mod(modulus);
//				// options[1][j] = options[0][j].multiply(message[j]).mod(modulus);
//
//				// Induce the delay
//				if(ModExp.delayMillis > 0)
//					try 
//					{
//						Thread.sleep(ModExp.delayMillis);
//					} 
//					catch (InterruptedException e) 
//					{
//						e.printStackTrace();
//						System.exit(-1);
//					}
//                end = time();
//
//				multTime = (end - start);
//				remainingTime[1][j] -= (multTime);					// this incorrect? choice yields
//				if( bitFileOutput ) p1.println((multTime));			// variability in the guesses
////				iterationTime[1][j] -= (multTime + squareTime);		// this choice makes all the guesses 0
////				if( bitFileOutput ) p1.println((multTime+squareTime));
//			}
//			
//			if( bitFileOutput )
//			{
//				p0.close();
//				p1.close();
//			}
//
//			// Calculate variances for time remaining after subtracting
//			// times for 0 and 1 bits
//			double var0 = Statistics.variance(remainingTime[0]);
//			double var1 = Statistics.variance(remainingTime[1]);
//
//			
//			int target, source;
//			
//			// If the Var(T) for bit=1 is smaller
//			// The first bit must be 1
//			if( (var0 > var1) || (i == 0) )
//			{
//				// bit is set
//				
//				// Set the running total to the candidate total for a 1-bit
//				// Set the candidate total for a 1-bit to the old running total
//				// (it will be overwritten)
//				BigInteger [] temp = runningTotal;
//				runningTotal = intermediateValue[1];
//				intermediateValue[1] = temp;
//
//				// Later copy the 1-bit remaining times to
//				// the 0-bit remaining times
//				target = 0;
//				source = 1;
//				
//				// Set the bit in the guess
//				guess = guess.setBit(exponent.bitLength()-i-1);
//				
//				System.out.print(1);
//			}
//			else
//			{
//				// bit not set
//
//				// Set the running total to the candidate total for a 0-bit
//				// Set the candidate total for a 0-bit to the old running total
//				// (it will be overwritten)
//				BigInteger [] temp = runningTotal;
//				runningTotal = intermediateValue[0];
//				intermediateValue[0] = temp;
//
//				// Later copy the 0-bit remaining times to
//				// the 1-bit remaining times
//				target = 1;
//				source = 0;
//
//				System.out.print(0);
//			}
//
//			// Output the variances
//			if (verbose) System.out.println("\t" + var0 + " " + var1);
//
//			// Transfer the remaining time values for the guessed bit
//			remainingTime[target] = Arrays.copyOf(remainingTime[source], remainingTime[source].length);
//		}
//		
//		return guess;
	}

}
